# -*- coding: utf-8 -*-
import numpy as np


def calculate_a_c(observed, forecast, threshold, probability):
    """calculate a for contingency table
    Args:
        observed is a 2d array from gauge;(times * gauge_num)
        forecast is a 3d array from simulation;(scheme_num * times * gauge_num)
        threshold is the user pre-defined threshold
        probability is of exceeding the selected threshold
    """
    height, row, column = forecast.shape
    observed = np.asarray(observed)
    forecast = np.asarray(forecast)

    forecast_gt_thd = forecast > threshold
    forecast_gt_thd_num = np.sum(forecast_gt_thd, axis=0)
    forecast_gt_thd_prob = forecast_gt_thd_num / height
    a = np.logical_and(observed > threshold,
                       forecast_gt_thd_prob > probability).sum()
    c = np.logical_and(observed > threshold,
                       forecast_gt_thd_prob <= probability).sum()
    return a, c


def calculate_b_d(observed, forecast, threshold, probability):
    """calculate a for contingency table
    Args:
        observed is a 2d array from gauge;(times * gauge_num)
        forecast is a 3d array from simulation;(scheme_num * times * gauge_num)
        threshold is the user pre-defined threshold
        probability is of exceeding the selected threshold
    """
    height, row, column = forecast.shape
    observed = np.asarray(observed)
    forecast = np.asarray(forecast)

    forecast_gt_thd = forecast > threshold
    forecast_gt_thd_num = np.sum(forecast_gt_thd, axis=0)
    forecast_gt_thd_prob = forecast_gt_thd_num / height

    b = np.logical_and(observed <= threshold,
                       forecast_gt_thd_prob > probability).sum()
    d = np.logical_and(observed <= threshold,
                       forecast_gt_thd_prob <= probability).sum()
    return b, d


def calculate_pod(observed, forecast, threshold, probability):
    try:
        a, c = calculate_a_c(observed, forecast, threshold, probability)
        numerator = a
        denominator = a + c
        numerator = numerator.astype(np.float64)
        denominator = denominator.astype(np.float64)
        # denominator[denominator == 0] = np.nan
        pod = numerator / denominator
        # pod_mean = np.mean(pod)
        return pod
    except Exception as e:
        raise e


def calculate_far(observed, forecast, threshold, probability):
    try:
        b, d = calculate_b_d(observed, forecast, threshold, probability)
        numerator = b
        denominator = b + d
        numerator = numerator.astype(np.float64)
        denominator = denominator.astype(np.float64)
        # denominator[denominator == 0] = np.nan
        far = numerator / denominator
        return far
    except Exception as e:
        raise e


def calculate_bias(observed, forecast, threshold, probability):
    try:
        a, c = calculate_a_c(observed, forecast, threshold, probability)
        b = calculate_b_d(observed, forecast, threshold, probability)[0]
        numerator = a + b
        denominator = a + c
        numerator = numerator.astype(np.float64)
        denominator = denominator.astype(np.float64)
        bias = numerator / denominator
        return bias
    except Exception as e:
        raise e


def calculate_roc_area(far, pod):
    area = []
    nums = len(far)
    for i in range(nums - 1):
        cur_area = (pod[i] + pod[i + 1]) * (far[i + 1] - far[i]) / 2
        area.append(cur_area)
    area = np.array(area)
    return np.sum(area)


def test(observed, forecast, threshold, probability):
    a, c = calculate_a_c(observed, forecast, threshold, probability)
    b, d = calculate_b_d(observed, forecast, threshold, probability)
    print(a + c + b + d)


if __name__ == '__main__':
    print("Calculating POD, FAR and Bias.")
